import matplotlib
matplotlib.use('TkAgg')  # 显式设置matplotlib后端为TkAgg

import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# 1. 加载Iris数据集
# 加载鸢尾花数据集，其中包含了花的特征数据以及对应的类别标签
iris_data = load_iris()
# 将特征数据转换为DataFrame格式，方便后续操作，并设置对应的列名
df_features = pd.DataFrame(data=iris_data.data, columns=iris_data.feature_names)
# 添加'species'列，存储实际的花的种类（类别标签）
df_features['species'] = iris_data.target

# 2. 数据标准化
# 创建标准化器对象，用于将数据特征缩放到均值为0，方差为1的范围
scaler = StandardScaler()
# 对除了最后一列（即类别标签列）的数据进行标准化转换
scaled_features = scaler.fit_transform(df_features[df_features.columns[:-1]])

# 3. 应用K-means算法
# 根据鸢尾花数据集中有三种不同花种类，设定聚类的簇数量为3
num_clusters = 3
# 创建KMeans聚类模型对象，设置聚类簇数，并指定随机种子以保证结果的可重复性
kmeans_model = KMeans(n_clusters=num_clusters, random_state=42)
# 对标准化后的特征数据进行聚类预测，并将聚类结果添加到DataFrame中作为新的一列'cluster'
df_features['cluster'] = kmeans_model.fit_predict(scaled_features)

# 4. 可视化聚类结果
# 创建一个新的图形对象，设置图形大小
fig = plt.figure(figsize=(10, 7))
# 在图形中添加一个3D坐标轴子图
ax = fig.add_subplot(111, projection='3d')

# 使用数据的前三个特征进行三维可视化，根据聚类结果设置点的颜色，使用viridis颜色映射，设置点的大小
scatter_plot = ax.scatter(df_features.iloc[:, 0], df_features.iloc[:, 1], df_features.iloc[:, 2],
                          c=df_features['cluster'], cmap='viridis', s=50)

# 添加颜色条用于展示不同聚类簇对应的颜色标识，并设置标题为'Clusters'
legend = ax.legend(*scatter_plot.legend_elements(), title="Clusters")
ax.add_artist(legend)

# 设置坐标轴的标签，分别对应花的不同特征名称
ax.set_xlabel('Sepal Length')
ax.set_ylabel('Sepal Width')
ax.set_zlabel('Petal Length')
# 设置图形的标题为'Iris数据集上的K-means聚类结果'
plt.title('K-means Clustering on Iris Dataset')
# 显示图形
plt.show()

# 5. 输出每个簇中样本的数量
print("Cluster distribution:")
print(df_features['cluster'].value_counts())

# 6. 比较实际类别与聚类结果（仅用于评估）
print("\nComparison with actual species:")
print(pd.crosstab(df_features['species'], df_features['cluster']))